import mujoco
import mujoco.viewer as viewer
from threading import Thread
import time
from Scripts.controller import Controller
from Scripts.sensor import Sensor

class Env():
    def __init__(
            self, xml_path: str, 
            enable_rangefinder_render: bool = True
        ):
        
        self.xml_path = xml_path
        self.enable_rangefinder_render = enable_rangefinder_render
        
        self.model = mujoco.MjModel.from_xml_path(self.xml_path)
        self.data = mujoco.MjData(self.model)
        
        self.sensor = Sensor(self.model, self.data)
        
        self.controller = Controller(self.model, self.data)
        self.target_vel = [0.0] * 2
        
    # 启动渲染
    def launch_viewer(self):
        self.viewer = viewer.launch_passive(self.model, self.data, show_left_ui=False, show_right_ui=False)
        if not self.enable_rangefinder_render:
            self.viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_RANGEFINDER] = False
        
        self.viewer_render_thread = Thread(target=self.viewer_render, daemon=True)
        self.viewer_render_thread.start()
        time.sleep(0.2)
        
    def viewer_render(self):
        while self.viewer.is_running():
            self.vel_control(self.target_vel)
            mujoco.mj_step(self.model, self.data)
            self.viewer.sync()
            time.sleep(0.002)
    
    # 关闭渲染
    def close_viewer(self):
        self.viewer.close()
    # 控制car速度
    def vel_control(self, target_vel: list):
        wheel_torque = self.controller.car_vel_controller(target_vel)
        self.data.ctrl[:4] = wheel_torque
    # 设置car目标速度
    def set_car_vel(self, target_vel):
        self.target_vel = target_vel